﻿using System;
using System.ComponentModel;
using System.Numerics;
using System.Security.Cryptography;

namespace SharpShare.Afp.Protocol.Security {
    [DisplayName("DHX2")]
    [Description("Diffie-Hellman Key Exchange 2.  This is the recommended UAM and is most secure.")]
    public class AfpDhx2UserAuthenticationMethod : IAfpUserAuthenticationMethod {
        public const int KeySize = 1024;

        private static readonly byte[] C2SIV = { 0x4c, 0x57, 0x61, 0x6c, 0x6c, 0x61, 0x63, 0x65 };
        private static readonly byte[] S2CIV = { 0x43, 0x4a, 0x61, 0x6c, 0x62, 0x65, 0x72, 0x74 };

        #region IAfpUserAuthenticationMethod Members

        public string Name {
            get { return AfpUams.kDHX2UAMStr; }
        }

        public AfpUserAuthenticationResult Authenticate(IAfpSession session, string version, string path, string username, AfpStream authInfoStream) {
            Context context = session.Properties.GetOrSet<Context>();

            context.Username = username;
            context.Version = version;

            int keyByteSize = (KeySize / 8);

            AfpStream stream = new AfpStream();
            stream.WriteBigInteger(context.G, 4);
            stream.WriteUInt16((ushort)keyByteSize);    // len
            stream.WriteBigInteger(context.P, keyByteSize);
            stream.WriteBigInteger(context.Mb, keyByteSize);

            return new AfpContinueUserAuthenticationResult(1, stream.ToByteArray());
        }

        public AfpUserAuthenticationResult ContinueAuthentication(IAfpSession session, short id, AfpStream authInfoStream) {
            Context context = session.Properties.Get<Context>();
            int keyByteSize = (KeySize / 8);

            switch (id) {
                case 1: {
                        context.Ma = authInfoStream.ReadBigInteger(keyByteSize);
                        byte[] encryptedClientNonce = authInfoStream.ReadBytes(16);
                        byte[] clientNonce = context.Decrypt(encryptedClientNonce, C2SIV);

                        context.ClientNonce = clientNonce.FromAfpNetworkByteArray();

                        // Create response
                        AfpStream stream = new AfpStream();

                        AfpStream responseEncryptStream = new AfpStream();
                        responseEncryptStream.WriteBigInteger(context.ClientNonce + 1, 16);
                        responseEncryptStream.WriteBigInteger(context.ServerNonce, 16);
                        byte[] responseDataToEncrypt = responseEncryptStream.ToByteArray();
                        byte[] responseEncryptedData = context.Encrypt(responseDataToEncrypt, S2CIV);

                        stream.WriteBytes(responseEncryptedData);

                        return new AfpContinueUserAuthenticationResult(2, stream.ToByteArray());
                    }
                case 2: {
                        byte[] encryptedNonceAndPassword = authInfoStream.ReadBytes(16 + 256); // 16 = Nonce size, 256 = password size
                        byte[] decryptedNonceAndPassword = context.Decrypt(encryptedNonceAndPassword, C2SIV);
                        AfpStream decryptedNoncePasswordStream = new AfpStream(decryptedNonceAndPassword);

                        BigInteger noncePlusOne = decryptedNoncePasswordStream.ReadBigInteger(16);

                        if (noncePlusOne != context.ServerNonce + 1) {
                            return new AfpDenyUserAuthenticationResult();
                        }

                        string password = System.Text.Encoding.ASCII.GetString(decryptedNoncePasswordStream.ReadBytes(256))
                            .TrimEnd('\0');

                        IAfpUser user = session.Server.UserManager.GetUser(context.Username);

                        if (user == null || !user.Authenticate(password))
                            return new AfpDenyUserAuthenticationResult();

                        return new AfpAllowUserAuthenticationResult(context.Version, this.Name, user);
                    }
                default:
                    throw new InvalidOperationException("Unexpected ID.");
            }
        }

        #endregion

        private sealed class Context {
            private const string PBase64 = "///////////JD9qiIWjCNMTGYouA3BzRKQJOCIpnzHQCC76mOxObIlFKCHmONATd75UZs806QxswKwpt8l8UN0/hNW1tUcJF5IW1dmJefsb0TELppjftawv/XLb0Brft7jhr+1qJn6WunyQRfEsf5kkoZlHs5lOB//////////8=";
            private const string GBase64 = "Fg==";

            private MD5Cng _md5 = new MD5Cng();
            private DotNetCrypt.CAST128Managed _cast = new DotNetCrypt.CAST128Managed();

            public Context() {
                _cast.ExtendedMode = DotNetCrypt.ExtendedCipherMode.CBC;
                _cast.Padding = PaddingMode.None;
                _cast.Mode = CipherMode.CBC;

                RNGCryptoServiceProvider rng = new RNGCryptoServiceProvider();

                byte[] rb = new byte[KeySize / 8];
                rng.GetBytes(rb);
                rb[rb.Length - 1] = 0;
                this.Rb = new BigInteger(rb);

                byte[] serverNonce = new byte[16];
                rng.GetBytes(serverNonce);
                serverNonce[serverNonce.Length - 1] = 0;
                this.ServerNonce = new BigInteger(serverNonce);
            }

            public string Username { get; set; }
            public string Version { get; set; }

            public BigInteger G {
                get {
                    byte[] useBytes = Convert.FromBase64String(GBase64);

                    if ((useBytes[useBytes.Length - 1] & 0x80) == 0x80) {
                        Array.Resize(ref useBytes, useBytes.Length + 1);
                    }

                    return new BigInteger(useBytes);
                }
            }
            public BigInteger P {
                get {
                    byte[] useBytes = Convert.FromBase64String(PBase64);

                    if ((useBytes[useBytes.Length - 1] & 0x80) == 0x80) {
                        Array.Resize(ref useBytes, useBytes.Length + 1);
                    }

                    return new BigInteger(useBytes);
                }
            }
            public BigInteger Rb { get; private set; }
            public BigInteger Mb {
                get {
                    return BigInteger.ModPow(this.G, this.Rb, this.P);
                }
            }
            public byte[] Key {
                get {
                    BigInteger result = BigInteger.ModPow(this.Ma, this.Rb, this.P);

                    // MD5 the result.
                    byte[] resultData = result.ToAfpNetworkByteArray(KeySize / 8);
                    byte[] hashed = _md5.ComputeHash(resultData);

                    return hashed;
                }
            }

            public BigInteger Ma { get; set; }

            public BigInteger ClientNonce { get; set; }

            public BigInteger ServerNonce { get; private set; }

            public byte[] Decrypt(byte[] data, byte[] ivData) {
                ICryptoTransform transform = _cast.CreateDecryptor(this.Key, ivData);
                byte[] decrypted = transform.TransformFinalBlock(data, 0, data.Length);
                return decrypted;
            }
            public byte[] Encrypt(byte[] data, byte[] ivData) {
                ICryptoTransform transform = _cast.CreateEncryptor(this.Key, ivData);
                byte[] encrypted = transform.TransformFinalBlock(data, 0, data.Length);
                return encrypted;
            }


        }

        /*
  1
 Client to server
 | FPLogin (2 bytes) | AFP Vers | ‘DHX2’ | Username (padded) |
 2
 Server to client
 | ID | g | len | p | Mb | and a result code
 3
 Client to server
 | FPLoginCont (2 bytes) | ID | Ma | (client nonce, C2SIV)K |
 4
 Server to client
 | ID + 1 | (clientNonce + 1, serverNonce, S2CIV)K | and a result code
 5
 Client to server
 | FPLoginCont (2 bytes) | ID + 1 | (serverNonce+1, password, C2SIV)K |
 6
 Server to client
 A result code of kFPNoErr if authentication was successful
  */
    }
}
